{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Using TPUs - Updated.ipynb",
      "private_outputs": true,
      "version": "0.3.2",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "accelerator": "TPU",
    "toc_visible": true
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Jxv6goXm7oGF"
      },
      "source": [
        "##### Copyright 2018 The TensorFlow Authors.\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "llMNufAK7nfK"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\"); { display-mode: \"form\" }\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "metadata": {
        "id": "Q-Y9SLvMl8IU",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/beta/{PATH}\">\n",
        "    <img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />\n",
        "    View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/{PATH}.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />\n",
        "    Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/{PATH}.ipynb\">\n",
        "    <img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />\n",
        "    View source on GitHub</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "pyxC7NBSdLda"
      },
      "cell_type": "markdown",
      "source": [
        "# Using TPUs\n",
        "\n",
        "Tensor Processing Units (TPUs) are Google's specialized ASICs designed to dramatically accelerate machine learning workloads. They are available on Google Colab, the TensorFlow Research Cloud and Google Compute Engine. \n",
        "\n",
        "In this notebook, you can try training a convolutional neural network against the Fashion MNIST dataset on Cloud TPUs using tf.keras and Distribution Strategy.\n"
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "EZMHru3ogQ8S"
      },
      "cell_type": "markdown",
      "source": [
        "## Learning Objectives\n",
        "\n",
        "In this Colab, you will learn how to:\n",
        "\n",
        "* Write a standard 4-layer conv-net with drop-out and batch normalization in Keras.\n",
        "* Use TPUs and Distribution Strategy to train the model.\n",
        "* Run a prediction to see how well the model can predict fashion categories and output the result."
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "KvRorTBAgiXE"
      },
      "cell_type": "markdown",
      "source": [
        "## Instructions\n",
        "\n",
        "To use TPUs in Colab:\n",
        "\n",
        "1. On the main menu, click Runtime and select **Change runtime type**. Set \"TPU\" as the hardware accelerator.\n",
        "1. Click Runtime again and select **Runtime > Run All**. You can also run the cells manually with Shift-ENTER. "
      ]
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "MpilHMdbjdWU"
      },
      "cell_type": "markdown",
      "source": [
        "## Data, Model, and Training\n",
        "\n",
        "### Download the Data\n",
        "\n",
        "Begin by downloading the fashion MNIST dataset using `tf.keras.datasets`, as shown below. We will also need to convert the data to `float32` format, as the data types supported by TPUs are limited right now.\n",
        "\n",
        "TPUs currently do not support Eager Execution, so we disable that with `disable_eager_execution()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "V9oECvVSI1Kj"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import, division, print_function, unicode_literals\n",
        "import numpy as np"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "bxzFVvMBjc_c",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "from __future__ import absolute_import, division, print_function\n",
        "\n",
        "!pip install tensorflow-gpu==2.0.0-beta1\n",
        "import tensorflow as tf\n",
        "tf.compat.v1.disable_eager_execution()\n",
        "\n",
        "import numpy as np\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = tf.keras.datasets.fashion_mnist.load_data()\n",
        "\n",
        "# add empty color dimension\n",
        "x_train = np.expand_dims(x_train, -1)\n",
        "x_test = np.expand_dims(x_test, -1)\n",
        "\n",
        "# convert types to float32\n",
        "x_train = x_train.astype(np.float32)\n",
        "x_test = x_test.astype(np.float32)\n",
        "y_train = y_train.astype(np.float32)\n",
        "y_test = y_test.astype(np.float32)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "doqxruT0pPcN"
      },
      "cell_type": "markdown",
      "source": [
        "### Initialize TPUStrategy\n",
        "\n",
        "We first initialize the TPUStrategy object before creating the model, so that Keras knows that we are creating a model for TPUs. \n",
        "\n",
        "To do this, we are first creating a TPUClusterResolver using the IP address of the TPU, and then creating a TPUStrategy object from the Cluster Resolver."
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "QSwckaHupcfv",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "import os\n",
        "\n",
        "resolver = tf.distribute.cluster_resolver.TPUClusterResolver()\n",
        "tf.tpu.experimental.initialize_tpu_system(resolver)\n",
        "strategy = tf.distribute.experimental.TPUStrategy(resolver)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "tqkpgoObohoj"
      },
      "cell_type": "markdown",
      "source": [
        "### Define the Model\n",
        "\n",
        "The following example uses a standard conv-net that has 4 layers with drop-out and batch normalization between each layer. Note that we are creating the model within a `strategy.scope`.\n"
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "q1Zk9I_2ol8M",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "with strategy.scope():\n",
        "  model = tf.keras.models.Sequential()\n",
        "  model.add(tf.keras.layers.BatchNormalization(input_shape=x_train.shape[1:]))\n",
        "  model.add(tf.keras.layers.Conv2D(64, (5, 5), padding='same', activation='elu'))\n",
        "  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))\n",
        "  model.add(tf.keras.layers.Dropout(0.25))\n",
        "\n",
        "  model.add(tf.keras.layers.BatchNormalization())\n",
        "  model.add(tf.keras.layers.Conv2D(128, (5, 5), padding='same', activation='elu'))\n",
        "  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2)))\n",
        "  model.add(tf.keras.layers.Dropout(0.25))\n",
        "\n",
        "  model.add(tf.keras.layers.BatchNormalization())\n",
        "  model.add(tf.keras.layers.Conv2D(256, (5, 5), padding='same', activation='elu'))\n",
        "  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))\n",
        "  model.add(tf.keras.layers.Dropout(0.25))\n",
        "\n",
        "  model.add(tf.keras.layers.BatchNormalization())\n",
        "  model.add(tf.keras.layers.Conv2D(512, (5, 5), padding='same', activation='elu'))\n",
        "  model.add(tf.keras.layers.MaxPooling2D(pool_size=(2, 2), strides=(2,2)))\n",
        "  model.add(tf.keras.layers.Dropout(0.25))\n",
        "\n",
        "  model.add(tf.keras.layers.Flatten())\n",
        "  model.add(tf.keras.layers.Dense(256))\n",
        "  model.add(tf.keras.layers.Activation('elu'))\n",
        "  model.add(tf.keras.layers.Dropout(0.5))\n",
        "  model.add(tf.keras.layers.Dense(10))\n",
        "  model.add(tf.keras.layers.Activation('softmax'))\n",
        "  model.summary()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "colab_type": "text",
        "id": "2pWxbBHIopia"
      },
      "cell_type": "markdown",
      "source": [
        "### Train on the TPU\n",
        "\n",
        "To train on the TPU, we can simply call `model.compile` under the strategy scope, and then call `model.fit` to start training. In this case, we are training for 5 epochs with 60 steps per epoch, and running evaluation at the end of 5 epochs.\n",
        "\n",
        "It may take a while for the training to start, as the data and model has to be transferred to the TPU and compiled before training can start."
      ]
    },
    {
      "metadata": {
        "colab_type": "code",
        "id": "3-prcoIQ2uZ-",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "with strategy.scope():\n",
        "  model.compile(\n",
        "      optimizer=tf.train.AdamOptimizer(learning_rate=1e-3),\n",
        "      loss=tf.keras.losses.sparse_categorical_crossentropy,\n",
        "      metrics=['sparse_categorical_accuracy']\n",
        "  )\n",
        "\n",
        "model.fit(\n",
        "    (x_train, y_train),\n",
        "    epochs=5,\n",
        "    steps_per_epoch=60,\n",
        "    validation_data=(x_test, y_test),\n",
        "    validation_freq=5,\n",
        ")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "fL1voAPek6jI",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### Check our results with Inference\n",
        "\n",
        "Now that we are done training, we can see how well the model can predict fashion categories:"
      ]
    },
    {
      "metadata": {
        "id": "c74exRcLk6jJ",
        "colab_type": "code",
        "colab": {}
      },
      "cell_type": "code",
      "source": [
        "LABEL_NAMES = ['t_shirt', 'trouser', 'pullover', 'dress', 'coat', 'sandal', 'shirt', 'sneaker', 'bag', 'ankle_boots']\n",
        "\n",
        "from matplotlib import pyplot\n",
        "%matplotlib inline\n",
        "\n",
        "def plot_predictions(images, predictions):\n",
        "  n = images.shape[0]\n",
        "  nc = int(np.ceil(n / 4))\n",
        "  f, axes = pyplot.subplots(nc, 4)\n",
        "  for i in range(nc * 4):\n",
        "    y = i // 4\n",
        "    x = i % 4\n",
        "    axes[x, y].axis('off')\n",
        "    \n",
        "    label = LABEL_NAMES[np.argmax(predictions[i])]\n",
        "    confidence = np.max(predictions[i])\n",
        "    if i > n:\n",
        "      continue\n",
        "    axes[x, y].imshow(images[i])\n",
        "    axes[x, y].text(0.5, -1.5, label + ': %.3f' % confidence, fontsize=12)\n",
        "\n",
        "  pyplot.gcf().set_size_inches(8, 8)  \n",
        "\n",
        "plot_predictions(np.squeeze(x_test[:16]), \n",
        "                 model.predict(x_test[:16]))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "metadata": {
        "id": "0E_gHZZQk6jN",
        "colab_type": "text"
      },
      "cell_type": "markdown",
      "source": [
        "### What's next\n",
        "\n",
        "* Learn about [Cloud TPUs](https://cloud.google.com/tpu/docs) that Google designed and optimized specifically to speed up and scale up ML workloads for training and inference and to enable ML engineers and researchers to iterate more quickly.\n",
        "* Explore the range of [Cloud TPU tutorials and Colabs](https://cloud.google.com/tpu/docs/tutorials) to find other examples that can be used when implementing your ML project.\n",
        "\n",
        "On Google Cloud Platform, in addition to GPUs and TPUs available on pre-configured [deep learning VMs](https://cloud.google.com/deep-learning-vm/),  you will find [AutoML](https://cloud.google.com/automl/)*(beta)* for training custom models without writing code and [Cloud ML Engine](https://cloud.google.com/ml-engine/docs/) which will allows you to run parallel trainings and hyperparameter tuning of your custom models on powerful distributed hardware."
      ]
    }
  ]
}